# -*- coding: utf-8 -*-
"""
Streamlit前端就绪检查脚本
检查所有依赖和配置是否正确
"""

import sys
from pathlib import Path

project_root = Path(__file__).resolve().parent
sys.path.insert(0, str(project_root))

def check_all():
    """检查所有依赖和配置"""
    
    print("=" * 60)
    print("Streamlit 前端就绪检查")
    print("=" * 60)
    
    all_ok = True
    
    # 1. 检查 Python 包
    print("\n【步骤 1】检查 Python 依赖包...")
    required_packages = {
        'streamlit': 'streamlit',
        'psycopg2': 'psycopg2-binary',
        'pandas': 'pandas',
        'openpyxl': 'openpyxl',
        'jieba': 'jieba',
        'dotenv': 'python-dotenv',
        'langchain': 'langchain',
        'langchain_openai': 'langchain-openai',
        'langchain_core': 'langchain-core'
    }
    
    missing_packages = []
    for package, install_name in required_packages.items():
        try:
            __import__(package)
            print(f"  ✓ {package}")
        except ImportError:
            print(f"  ✗ {package} (缺失)")
            missing_packages.append(install_name)
            all_ok = False
    
    if missing_packages:
        print(f"\n  请安装缺失的包:")
        print(f"  pip install {' '.join(missing_packages)}")
    
    # 2. 检查数据库连接
    print("\n【步骤 2】检查数据库连接...")
    try:
        import psycopg2
        from D_db.core.config import DB_CONFIG
        
        import os
        os.environ['PGCLIENTENCODING'] = 'UTF8'
        
        conn = psycopg2.connect(**DB_CONFIG)
        cursor = conn.cursor()
        
        # 检查表是否存在
        cursor.execute("""
            SELECT EXISTS (
                SELECT FROM information_schema.tables 
                WHERE table_name = 'qa_knowledge'
            );
        """)
        table_exists = cursor.fetchone()[0]
        
        if table_exists:
            # 检查数据量
            cursor.execute("SELECT COUNT(*) FROM qa_knowledge;")
            count = cursor.fetchone()[0]
            print(f"  ✓ 数据库连接成功")
            print(f"  ✓ qa_knowledge 表存在")
            print(f"  ✓ 数据量: {count} 条")
        else:
            print(f"  ✗ qa_knowledge 表不存在")
            print(f"  请先运行: python setup_db_and_import.py")
            all_ok = False
        
        cursor.close()
        conn.close()
        
    except Exception as e:
        print(f"  ✗ 数据库连接失败: {e}")
        print(f"  请检查 PostgreSQL 服务和配置")
        all_ok = False
    
    # 3. 检查 .env 文件
    print("\n【步骤 3】检查环境配置...")
    from dotenv import load_dotenv
    import os
    
    env_file = project_root / '.env'
    if not env_file.exists():
        print(f"  ✗ .env 文件不存在")
        print(f"  请创建 .env 文件并配置 DASHSCOPE_API_KEY")
        all_ok = False
    else:
        load_dotenv()
        api_key = os.getenv('DASHSCOPE_API_KEY')
        base_url = os.getenv('BASE_URL')
        
        if api_key and api_key != 'your_api_key_here':
            print(f"  ✓ .env 文件存在")
            print(f"  ✓ DASHSCOPE_API_KEY 已配置")
            print(f"  ✓ BASE_URL: {base_url}")
        else:
            print(f"  ✗ DASHSCOPE_API_KEY 未配置或无效")
            print(f"  请在 .env 文件中设置有效的 API Key")
            all_ok = False
    
    # 4. 检查核心模块
    print("\n【步骤 4】检查核心模块...")
    try:
        from D_db.core.db_manager import DBManager
        print(f"  ✓ DBManager")
        
        from D_db.core.tokenizer import Tokenizer
        print(f"  ✓ Tokenizer")
        
        from D_db.tools.prompt_utils import build_llm_context_from_qa
        print(f"  ✓ prompt_utils")
        
        from D_db.core.config import RESERVOIR_MAPPING
        print(f"  ✓ RESERVOIR_MAPPING ({len(RESERVOIR_MAPPING)} 个水库)")
        
    except Exception as e:
        print(f"  ✗ 核心模块导入失败: {e}")
        all_ok = False
    
    # 5. 检查 Streamlit 应用文件
    print("\n【步骤 5】检查 Streamlit 应用文件...")
    app_files = [
        'app.py',
        'run_streamlit.bat'
    ]
    
    for file in app_files:
        file_path = project_root / file
        if file_path.exists():
            print(f"  ✓ {file}")
        else:
            print(f"  ✗ {file} (缺失)")
            all_ok = False
    
    # 总结
    print("\n" + "=" * 60)
    if all_ok:
        print("✓ 所有检查通过！系统已就绪")
        print("=" * 60)
        print("\n可以启动 Streamlit 应用了：")
        print("\n方法 1 (Windows): 双击 run_streamlit.bat")
        print("方法 2 (命令行): streamlit run app.py")
        print("\n应用将在浏览器中打开: http://localhost:8501")
    else:
        print("✗ 部分检查未通过，请先解决上述问题")
        print("=" * 60)
    
    return all_ok

def main():
    try:
        result = check_all()
        sys.exit(0 if result else 1)
    except Exception as e:
        print(f"\n检查过程出错: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    main()

